from langchain.chains.conversation.base import ConversationChain
from langchain.chains.llm import LLMChain
from langchain.chains.router import LLMRouterChain, MultiPromptChain
from langchain.chains.router.llm_router import RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAI

import __init__

# 物理链
physics_template = """您是一位非常聪明的物理教授。\n
您擅长以简洁易懂的方式回答物理问题。\n
当您不知道问题答案的时候，您会坦率承认不知道。\n
下面是一个问题：
{input}"""
physics_prompt = PromptTemplate.from_template(physics_template)

# 数学链
math_template = """您是一位非常优秀的数学教授。\n
您擅长回答数学问题。\n
您之所以如此优秀，是因为您能够将困难问题分解成组成的部分，回答这些部分，然后将它们组合起来，回答更广泛的问题。\n
{input}"""
math_prompt = PromptTemplate.from_template(math_template)

prompt_infos = [
    {
        "name": "physics",
        "description": "擅长回答物理问题",
        "prompt_template": physics_template,
    },
    {
        "name": "math",
        "description": "删除回答数学问题",
        "prompt_template": math_template,
    }
]

llm = OpenAI(
    temperature=0,
    model="gpt-3.5-turbo-instruct"
)
destination_chains = {}
for p_info in prompt_infos:
    name = p_info["name"]
    prompt_template = p_info["prompt_template"]
    prompt = PromptTemplate(
        template=prompt_template,
        input_variables=["input"]
    )
    chain = LLMChain(
        llm=llm,
        prompt=prompt,
    )
    destination_chains[name] = chain
default_chain = ConversationChain(
    llm=llm,
    output_key="text"
)

destinations = [f"{p['name']}:{p['description']}" for p in prompt_infos]
destinations_str = "\n".join(destinations)
router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)
print(MULTI_PROMPT_ROUTER_TEMPLATE)

router_prompt = PromptTemplate(
    template=router_template,
    input_variables=["input"],
    output_parser=RouterOutputParser()
)
router_chain = LLMRouterChain.from_llm(
    llm,
    router_prompt
)

chain = MultiPromptChain(
    router_chain=router_chain,
    destination_chains=destination_chains,
    default_chain=default_chain,
    verbose=True
)

print(chain.run("什么是牛顿第一定律？"))

print(chain.run("2+2等于几？"))

print(chain.run("两个黄鹂鸣翠柳，下一句？"))
